# -*- coding: utf-8 -*-
"""
-------------------------------------------------
File Name： img_augment
Description :
Author : 'li'
date： 2022/6/15
Change Activity:
2022/6/15:
-------------------------------------------------
"""
import random

from ml.common.design_pattern.pipline.base_module import BasePiplineModule
from ml.dl.misc.img_augment.flip import horizontal_flip


class RandomHorizontalFlip(BasePiplineModule):
    """
    flip image
    """

    def __init__(self, input_key='img', output_key='img', target_input_key='targets', target_output_key='targets', p=0.5):
        """

        Args:
            input_key:
            output_key:
            target_input_key:
            target_output_key:
            p:
        """
        super().__init__(input_key, output_key)
        self.p = p
        self.target_input_key = target_input_key
        self.target_output_key = target_output_key

    def __call__(self, data):
        img = data[self.input_key].copy()
        target = data[self.target_input_key].copy()
        if random.random() < self.p:
            flipped_image, target = horizontal_flip(img, target)
        data[self.output_key] = img
        data[self.target_output_key] = target
        return data
